﻿Imports System.Net.Sockets

''' <summary>
''' Provides simple, event-driven TCP client/server functionality.
''' </summary>
''' <remarks></remarks>
<System.ComponentModel.DefaultEvent("DataReceived")>
Public Class Winsock
    Inherits System.ComponentModel.Component

    ''' <summary>
    ''' Occurs when the Winsock component recieves data from a remote host for processing.
    ''' This event represents client requests for a Winsock component in server mode, and
    ''' server responses for a Winsock component in client mode.
    ''' </summary>
    ''' <remarks></remarks>
    Public Event DataReceived As EventHandler(Of DataReceivedEventArgs)
    ''' <summary>
    ''' Occurs when the Winsock component experiences an error.  Examine the exception,
    ''' inner exception, and WinsockErrorCategory to determine the cause of the error.
    ''' </summary>
    ''' <remarks></remarks>
    Public Event WinsockError As EventHandler(Of WinsockErrorEventArgs)

    Protected _RemoteClients As New System.Collections.Concurrent.ConcurrentBag(Of ClientInfo)
    Protected _ReceivedData As New System.Collections.Concurrent.ConcurrentQueue(Of ReceivedData)

    Private _ContainerControl As ContainerControl
    Private _ServerListener As TcpListener
    Private _LocalClient As ClientInfo

    ''' <summary>
    ''' Gets or sets the parent control which hosts this component and is responsible for the UI thread.
    ''' </summary>
    ''' <value>The container to be parent of this component.</value>
    ''' <returns>The parent container of this component.</returns>
    ''' <remarks></remarks>
    Public Property ContainerControl() As ContainerControl
        Get
            Return _ContainerControl
        End Get
        Set(ByVal value As ContainerControl)
            _ContainerControl = value
        End Set
    End Property

    Public Property LocalIpAddress As System.Net.IPAddress = System.Net.IPAddress.Any

    Private _Mode As WinsockMode
    ''' <summary>
    ''' Gets the current operating mode of the Winsock component.
    ''' </summary>
    ''' <value>A valid WinsockMode enumeration value.</value>
    ''' <returns>The current operating mode.</returns>
    ''' <remarks></remarks>
    Public ReadOnly Property Mode As WinsockMode
        Get
            Return _Mode
        End Get
    End Property

    ''' <summary>
    ''' Gets the currently connected port number, or sets the port number to connect to or listen on.
    ''' </summary>
    ''' <value>The port number to connect to or listen on.</value>
    ''' <returns>The currently connected port number or last specified port.</returns>
    ''' <remarks></remarks>
    Public Property Port As Integer

    Public Overrides Property Site() As System.ComponentModel.ISite
        Get
            Return MyBase.Site
        End Get
        Set(ByVal value As System.ComponentModel.ISite)
            MyBase.Site = value
            If value IsNot Nothing Then
                Dim host As System.ComponentModel.Design.IDesignerHost
                host = value.GetService(GetType(System.ComponentModel.Design.IDesignerHost))
                If host IsNot Nothing Then
                    _ContainerControl = CType(host.RootComponent, ContainerControl)
                Else
                    _ContainerControl = Nothing
                End If
            End If
        End Set
    End Property

    Private _IsRunning As Boolean
    ''' <summary>
    ''' Determines if the Winsock component is currently operating.
    ''' </summary>
    ''' <returns>True if the Winsock component is operating, otherwise false.</returns>
    ''' <remarks></remarks>
    Public ReadOnly Property IsRunning As Boolean
        Get
            Return _IsRunning
        End Get
    End Property

    ''' <summary>
    ''' Ends the current active mode of the Winsock component.
    ''' </summary>
    ''' <remarks></remarks>
    Public Sub Close()
        If _IsRunning Then
            If _Mode = WinsockMode.Client Then
                _IsRunning = False
                _LocalClient.Client.Close()
                _LocalClient = Nothing
                _Mode = WinsockMode.Closed
            ElseIf _Mode = WinsockMode.Server Then
                _IsRunning = False
                _ServerListener.Stop()
                _ServerListener = Nothing
                _Mode = WinsockMode.Closed
            End If
        End If
    End Sub

    ''' <summary>
    ''' Starts the Winsock component as a TCP client, using the port number specified by Port property.
    ''' </summary>
    ''' <param name="address"></param>
    ''' <remarks></remarks>
    Public Sub Connect(address As System.Net.IPAddress)
        Connect(address, Port)
    End Sub

    ''' <summary>
    ''' Starts the Winsock component as a TCP client.
    ''' </summary>
    ''' <param name="address">The IP address of the server to connect with.</param>
    ''' <param name="port">The port to connect to on the server.</param>
    ''' <remarks></remarks>
    Public Sub Connect(address As System.Net.IPAddress, port As Integer)
        If Not _IsRunning Then
            Dim client As TcpClient = New TcpClient()
            Try
                client.Connect(address, port)
                _Port = port
                _IsRunning = True
                _Mode = WinsockMode.Client
                _LocalClient = New ClientInfo(client)
                Try
                    DoBeginReadData(_LocalClient)
                    System.Threading.Tasks.Task.Factory.StartNew(AddressOf DoProcessReceivedData)
                Catch ex As Exception
                    ProcessWinsockError(New Exception("Error while attempting to wait for data on new connection. See inner exception for details.", ex), WinsockErrorCategory.Connection)
                    _IsRunning = False
                    _Mode = WinsockMode.Closed
                    _LocalClient = Nothing
                End Try
            Catch ex As Exception
                ProcessWinsockError(New Exception("Error connecting to server. See inner exception for details.", ex), WinsockErrorCategory.Connection)
            End Try
        Else
            ProcessWinsockError(New InvalidOperationException("Winsock component is already started."), WinsockErrorCategory.Connection)
        End If
    End Sub

    ''' <summary>
    ''' Starts the Winsock component as a TCP server, listening on the port number specified by Port property.
    ''' </summary>
    ''' <remarks></remarks>
    Public Sub Listen()
        Listen(Port)
    End Sub

    ''' <summary>
    ''' Starts the Winsock component as a TCP server.
    ''' </summary>
    ''' <param name="port">The port on which to listen for incomming connections.</param>
    ''' <remarks></remarks>
    Public Sub Listen(port As Integer)
        If Not _IsRunning Then
            Try
                _ServerListener = New TcpListener(LocalIpAddress, port)
                _ServerListener.Start()
                _Port = port
                _IsRunning = True
                _Mode = WinsockMode.Server
                Try
                    DoBeginConnectRemoteClient(_ServerListener)
                    System.Threading.Tasks.Task.Factory.StartNew(AddressOf DoProcessReceivedData)
                Catch ex As Exception
                    ProcessWinsockError(New Exception("Error while attempting to listen for client connection. See inner exception for details.", ex), WinsockErrorCategory.Connection)
                End Try
            Catch ex As Exception
                ProcessWinsockError(New Exception("Error creating server. See inner exception for details.", ex), WinsockErrorCategory.Connection)
            End Try
        Else
            ProcessWinsockError(New InvalidOperationException("Winsock component is already started."), WinsockErrorCategory.Connection)
        End If
    End Sub

    ''' <summary>
    ''' Submits a Winsock error to the component's error handler. Calling this method will subsequently raise
    ''' the WinsockError event.
    ''' </summary>
    ''' <param name="ex">The exception to submit.</param>
    ''' <param name="category">A valid WinsockErrorCategory enumeration value.</param>
    ''' <remarks></remarks>
    Public Sub ProcessWinsockError(ex As Exception, category As WinsockErrorCategory)
        Dim args As New WinsockErrorEventArgs(ex, category)
        DoRaiseEvent(New OnDataErrorDelegate(AddressOf OnDataError), args)
    End Sub

    ''' <summary>
    ''' (Client-Only) Attempts to write a block of data from the client to the server.
    ''' This method throws an InvalidOperationException if the Winsock component is in Server mode when called.
    ''' </summary>
    ''' <param name="data">The byte data to write to the server.</param>
    ''' <remarks></remarks>
    Public Sub Write(data() As Byte)
        If _Mode = WinsockMode.Client Then
            Try
                _LocalClient.Stream.Write(data, 0, data.Length)
            Catch ex As Exception
                ProcessWinsockError(New Exception("Error writing data. See inner exception for details.", ex), WinsockErrorCategory.Data)
            End Try
        Else
            ProcessWinsockError(New InvalidOperationException("Invalid Mode; must be client."), WinsockErrorCategory.Data)
        End If
    End Sub

    Protected Sub DoBeginConnectRemoteClient(listener As TcpListener)
        Try
            Dim info As New ConnectionInfo(listener)
            info.AsyncResult = listener.BeginAcceptTcpClient(AddressOf DoAcceptTcpClient, info)
        Catch ex As Exception
            ProcessWinsockError(New Exception("Error while attempting to listen for client connection. See inner exception for details.", ex), WinsockErrorCategory.Connection)
        End Try
    End Sub

    Protected Sub DoAcceptTcpClient(asyncResult As IAsyncResult)
        If _IsRunning Then
            Dim connection As ConnectionInfo = CType(asyncResult.AsyncState, ConnectionInfo)
            Try
                Dim client As TcpClient = connection.Listener.EndAcceptTcpClient(asyncResult)
                Dim info As New ClientInfo(client)
                _RemoteClients.Add(info)
                DoBeginReadData(info)
            Catch ex As Exception
                ProcessWinsockError(New Exception("Error while attempting to accept client connection.", ex), WinsockErrorCategory.Connection)
            End Try
            DoBeginConnectRemoteClient(connection.Listener)
        End If
    End Sub

    Protected Sub DoBeginReadData(info As ClientInfo)
        info.AsyncResult = info.Stream.BeginRead(info.ReadBuffer, 0, info.ReadBufferLength, AddressOf DoEndReadData, info)
    End Sub

    Protected Sub DoEndReadData(asyncResult As IAsyncResult)
        If _IsRunning Then
            Dim info As ClientInfo = CType(asyncResult.AsyncState, ClientInfo)
            Try
                Dim dataLength As Integer = info.Stream.EndRead(asyncResult)
                If dataLength > 0 Then
                    Dim data() As Byte = info.ReadBuffer.Take(dataLength).ToArray
                    _ReceivedData.Enqueue(New ReceivedData(info, data))
                End If
                DoBeginReadData(info)
            Catch ex As Exception
                ProcessWinsockError(New Exception("Connection terminated while reading data.", ex), WinsockErrorCategory.Data)
            End Try
        End If
    End Sub

    Protected Sub DoProcessReceivedData()
        Do
            If _ReceivedData.Count > 0 Then
                Try
                    Dim entry As ReceivedData = Nothing
                    If _ReceivedData.TryDequeue(entry) Then
                        Dim args As New DataReceivedEventArgs(entry)
                        DoRaiseEvent(New OnDataReceivedDelegate(AddressOf OnDataReceived), args)
                    End If
                Catch ex As Exception
                    ProcessWinsockError(New Exception("Error while processing received data. See inner exception for details.", ex), WinsockErrorCategory.Processing)
                End Try
            End If
            System.Threading.Thread.Sleep(0)
        Loop While _IsRunning
    End Sub

    Protected Sub DoRaiseEvent(method As [Delegate], ParamArray arguments() As Object)
        If _ContainerControl IsNot Nothing AndAlso _ContainerControl.InvokeRequired Then
            _ContainerControl.Invoke(method, arguments)
        Else
            method.Method.Invoke(Me, arguments)
        End If
    End Sub

    Protected Delegate Sub OnDataErrorDelegate(e As WinsockErrorEventArgs)
    Protected Sub OnDataError(e As WinsockErrorEventArgs)
        RaiseEvent WinsockError(Me, e)
    End Sub

    Protected Delegate Sub OnDataReceivedDelegate(e As DataReceivedEventArgs)
    Protected Sub OnDataReceived(e As DataReceivedEventArgs)
        RaiseEvent DataReceived(Me, e)
    End Sub

#Region " Windows Forms Designer Support "
    <System.Diagnostics.DebuggerNonUserCode()> _
    Public Sub New(ByVal container As System.ComponentModel.IContainer)
        MyClass.New()
        If (container IsNot Nothing) Then
            container.Add(Me)
        End If
    End Sub

    <System.Diagnostics.DebuggerNonUserCode()> _
    Public Sub New()
        MyBase.New()
        InitializeComponent()
    End Sub

    <System.Diagnostics.DebuggerNonUserCode()> _
    Protected Overrides Sub Dispose(ByVal disposing As Boolean)
        Try
            If disposing AndAlso components IsNot Nothing Then
                components.Dispose()
            End If
        Finally
            MyBase.Dispose(disposing)
        End Try
    End Sub

    Private components As System.ComponentModel.IContainer
    <System.Diagnostics.DebuggerStepThrough()> _
    Private Sub InitializeComponent()
        components = New System.ComponentModel.Container()
    End Sub
#End Region
End Class
